package leetcode;

import java.util.*;

/**
 * 1584. 连接所有点的最小费用     https://leetcode.cn/problems/min-cost-to-connect-all-points/
 */
public class LeetCode1584 {

    static class Graph {
        public Map<Integer, Node> nodes;
        public Set<Edge> edges;

        public Graph() {
            this.nodes = new HashMap<>();
            this.edges = new HashSet<>();
        }
    }

    static class Node {
        public int value;
        public int x;
        public int y;
        public ArrayList<Node> nexts;
        public ArrayList<Edge> edges;

        public Node(int value, int x, int y) {
            this.value = value;
            this.x = x;
            this.y = y;
            this.nexts = new ArrayList<>();
            this.edges = new ArrayList<>();
        }
    }

    static class Edge {
        public int weight;
        public Node from;
        public Node to;

        public Edge(Node from, Node to) {
            this.from = from;
            this.to = to;
            this.computeWight();
        }

        public void computeWight() {
            this.weight = Math.abs(this.from.x - this.to.x) + Math.abs(this.from.y - this.to.y);
        }
    }

    static class UnionFind {
        public Map<Node, Node> parent;
        public Map<Node, Integer> size;

        public UnionFind() {
            this.parent = new HashMap<>();
            this.size = new HashMap<>();
        }

        public void init(ArrayList<Node> values) {
            parent.clear();
            size.clear();
            for (Node value : values) {
                parent.put(value, value);
                size.put(value, 1);
            }
        }

        public Node find(Node cur) {
            Stack<Node> stack = new Stack<>();
            while (cur != parent.get(cur)) {
                stack.add(cur);
                cur = parent.get(cur);
            }
            while (!stack.isEmpty()) {
                parent.put(stack.pop(), cur);
            }
            return cur;
        }

        public void union(Node a, Node b) {
            if (a == null || b == null) {
                return;
            }
            Node node1 = find(a);
            Node node2 = find(b);
            if (node1 != node2) {
                if (size.get(node1) >= size.get(node2)) {
                    parent.put(node2, node1);
                    size.put(node1, size.get(node1) + size.get(node2));
                    size.remove(node2);
                } else {
                    parent.put(node1, node2);
                    size.put(node2, size.get(node1) + size.get(node2));
                    size.remove(node1);
                }
            }
        }

        public boolean isSameSet(Node node1, Node node2) {
            return find(node1) == find(node2);
        }
    }

    public static int minCostConnectPoints(int[][] points) {
        Graph graph = new Graph();
        Map<Integer, Node> nodes = graph.nodes;
        Set<Edge> edges = graph.edges;
        // 小根堆获取最小的边，用于每次弹出最小权重边
        PriorityQueue<Edge> heap = new PriorityQueue<>((o1, o2) -> o1.weight - o2.weight);
        for (int i = 0; i < points.length; i++) {
            nodes.put(i, new Node(i, points[i][0], points[i][1]));
        }
        for (int i = 0; i < points.length; i++) {
            Node nodeI = nodes.get(i);
            // 构建点、边
            for (int j = i + 1; j < points.length; j++) {
                Node nodeJ = nodes.get(j);
                nodeI.nexts.add(nodeJ);
                Edge edge = new Edge(nodeI, nodeJ);
                edges.add(edge);
                heap.add(edge);
                nodeI.edges.add(edge);
            }
        }
        // 构建并查集
        UnionFind unionFind = new UnionFind();
        unionFind.init(new ArrayList<Node>(nodes.values()));
        int ans = 0;
        while (!heap.isEmpty()) {
            Edge edge = heap.poll();
            if (!unionFind.isSameSet(edge.from, edge.to)) {
                ans += edge.weight;
                unionFind.union(edge.from, edge.to);
            }
        }
        return ans;
    }

    public static void main(String[] args) {
        int[][] points = {{0, 0}, {2, 2}, {3, 10}, {5, 2}, {7, 0}};
        System.out.println(minCostConnectPoints(points));
    }
}
